Skip to content

Make get_differential_vars type stable #2698

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 5 commits into
base: master
Choose a base branch
from

Conversation

Ickaser
Copy link
Contributor

@Ickaser Ickaser commented May 5, 2025

Checklist

  • Appropriate tests were added
  • Any code changes were done in a way that does not break public API
  • All documentation related to code changes were updated
  • The new code follows the
    contributor guidelines, in particular the SciML Style Guide and
    COLPRAC.
  • Any new documentation only uses public API

Additional context

Solves #2594 , based on #2594 (comment).

Question on testing

It seems like a good idea to add a test that would have caught #2594; my attempt would be to add @inferred in a couple places. Would

@testset "Inplace: $(isinplace(_prob)), DAEProblem: $(_prob isa DAEProblem), BrownBasic: $(initalg isa BrownFullBasicInit), Autodiff: $autodiff" for _prob in [
prob_mm, prob_mm_oop],
initalg in [BrownFullBasicInit(), ShampineCollocationInit()], autodiff in [true, false]
alg = Rodas5P(; autodiff)
function f(p)
sol = solve(remake(_prob, p = p), alg, abstol = 1e-14,
reltol = 1e-14, initializealg = initalg)
sum(sol)
end
@test ForwardDiff.gradient(f, [0.04, 3e7, 1e4])[0, 0, 0] atol=1e-8
end

or
p = [0.04, 3e7, 1e4]
u₀ = [1.0, 0, 0]
du₀ = [-0.04, 0.04, 0.0]
tspan = (0.0, 100000.0)
differential_vars = [true, true, false]
prob = DAEProblem(f, du₀, u₀, tspan, p, differential_vars = differential_vars)
prob_oop = DAEProblem{false}(f, du₀, u₀, tspan, p, differential_vars = differential_vars)
sol1 = solve(prob, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
sol2 = solve(prob_oop, DFBDF(), dt = 1e-5, abstol = 1e-8, reltol = 1e-8)
# These tests flex differentiation of the solver and through the initialization
# To only test the solver part and isolate potential issues, set the initialization to consistent
@testset "Inplace: $(isinplace(_prob)), DAEProblem: $(_prob isa DAEProblem), BrownBasic: $(initalg isa BrownFullBasicInit), Autodiff: $autodiff" for _prob in [
prob, prob_oop],
initalg in [BrownFullBasicInit(), ShampineCollocationInit()], autodiff in [true, false]
alg = DFBDF(; autodiff)
function f(p)
sol = solve(remake(_prob, p = p), alg, abstol = 1e-14,
reltol = 1e-14, initializealg = initalg)
sum(sol)
end
@test ForwardDiff.gradient(f, [0.04, 3e7, 1e4])[0, 0, 0] atol=1e-8
end

be a good place for these? Should this get tested with all the implicit solvers where the mass-matrix DAE formulation is allowed?

@ChrisRackauckas
Copy link
Member

be a good place for these?

Yeah, the more the marrier.

https://github.com/SciML/OrdinaryDiffEq.jl/blob/master/test/interface/mass_matrix_tests.jl is probably a good test to slam some around.

@Ickaser
Copy link
Contributor Author

Ickaser commented May 5, 2025

Turns out that lots of these are still not inferrable--even with this PR, going from Rosenbrock DAE AD tests, the move from

sol = solve(prob_mm, Rodas5P(), reltol = 1e-8, abstol = 1e-8)

to

using ADTypes: AutoForwardDiff
sol = @inferred solve(prob_mm, Rodas5P(autodiff=AutoForwardDiff(chunksize=3)), reltol = 1e-8, abstol = 1e-8)

already runs into an issue, since the type of solve is still inferred to be Any. I think there is something else going on here (or I don't know how to use @inferred properly, which is possible), but I can't tell what makes this different from the MWE I've been testing against.

That said, even though this PR apparently doesn't do enough for the existing test suite, it is already helping with the MWE and therefore probably my code base.

@ChrisRackauckas
Copy link
Member

It's fine to sprinkle some @test_brokens around to get things in, and document them.

@Ickaser
Copy link
Contributor Author

Ickaser commented May 6, 2025

OK, so the quirk in the test I was looking at turns out to be that for

M = Diagonal([1.0, 1.0, 0.0])
roberf = ODEFunction(rober, mass_matrix = M)

get_differential_vars(roberf, u0) is fully inferrable, since M was constructed as diagonal. But in the test, we had

M = [1.0 0 0
     0 1.0 0
     0 0 0]
roberf = ODEFunction(rober, mass_matrix = M)

With a full matrix, inference's best guess at get_differential_vars(roberf, u0) is Union{BitVector, DifferentialVarsUndefined}.

We currently get DifferentialVarsUndefined if all of the following are satisfied:

  • Matrix is not a UniformScaling
  • Matrix has at least one zero element
  • Matrix is either an AbstractSciMLOperator or not diagonal

So, for example, a matrix like [1 1 0; 0 1 0; 0 0 1] with one off-diagonal nonzero currently gives DifferentialVarsUndefined.
What should be the exact criteria for having the differential vars undefined? My instinct is that even a fully-zero matrix should still just give us a falses(size(mass_matrix, 1)), so maybe this branch should not even exist?

@ChrisRackauckas
Copy link
Member

What should be the exact criteria for having the differential vars undefined? My instinct is that even a fully-zero matrix should still just give us a falses(size(mass_matrix, 1)), so maybe this branch should not even exist?

I believe the proper definition is that the number of algebraic variables is determined by the size of the null space of M. If M is not diagonal the algebraic variables are the e_i vectors which live in the null space? I think you can always orthogonalize a basis of the null space down to just being defined by linear combinations the algebraic variables.

But I couldn't figure out how to calculate this easily so the diagonal cases were handled and any case where someone couldn't figure out the differential variables but needed it turned into an error. In lots of cases it's not necessary to calculate so it ended up being an okay solution, with a few special cases like dense. But yes... I don't know, I think in general you probably need to QR factorize the M and then take the Q part and check its null space?

julia> M = [1 1 0; 1 1 0; 0 0 0]
3×3 Matrix{Int64}:
 1  1  0
 1  1  0
 0  0  0

julia> qr(M).R
3×3 Matrix{Float64}:
 -1.41421  -1.41421  0.0
  0.0       0.0      0.0
  0.0       0.0      0.0

That definition isn't quite right either.

@Ickaser
Copy link
Contributor Author

Ickaser commented May 6, 2025

But I couldn't figure out how to calculate this easily so the diagonal cases were handled and any case where someone couldn't figure out the differential variables but needed it turned into an error. In lots of cases it's not necessary to calculate so it ended up being an okay solution, with a few special cases like dense.

OK, that makes sense--I had the sense my solution might be too easy, otherwise it would have already been there. I'll revert it to keep the previous behavior, I think

For at least some of the tests, I can just change to explicitly using a Diagonal matrix, and do likewise in the examples to point people this way, plus maybe add a note on this type-stability question to the mass-matrix docs. That would be easy enough, and that way we avoid doing any possibly-expensive linear algebra (i.e., for cases with larger mass matrices) in what should typically be a simple preparation step.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants